from matplotlib.finance import quotes_historical_yahoo
from matplotlib.dates import num2date
import datetime
import urllib2
import numpy
import pandas

class StockData:
    """
    Fetch stock data. The class implements methods able to acquire data from Yahoo Finance about companies within the specified time frame. Optionally you can download the list of companies from `here <http://www.student.dtu.dk/~s110848/companylist.csv>`_ and perform the actions for chosen companies, that are related to software sector. 
    
    .. note:: An instance can be created also without using any arguments.
	
    Arguments:
	
    * company_list -- list of stocknames for which the data will be fetched,
    * start_date -- the date that specifies the beginning of the timespan from which the stockdata will be fetched,
    * end_date -- the date that specifies the end of the timespan from which the stockdata will be fetched, the end date can be input as the number of days we would like to fetch starting from the start date.
	
    Sample usage of the class:
	
    >>> import datetime
    >>> import stockdata
    >>> START = datetime.date(2013, 1, 10)
    >>> NUM_DAYS = 10
    >>> sd = stockdata.StockData()
    >>> DATA = sd.get_stock_data(["MSFT", "GOOG"], START, NUM_DAYS)
    >>> DATA.MSFT[0]
    26.238379272326345
    >>> DATA.GOOG[1]
    737.0
    >> DATA.GOOG[4]
    742.0
    
    It consist of class initialization as well as fetching the stock prices within the date range 10-01-2013 to 20-01-2013. Below you can the the methods that this package provides in alphabetical order.
	
    """
    def __init__(self, company_list=None, start_date=None, end_date=None):
#        try:
        self.start_date = start_date
        self.end_date = end_date
        self.stock_data = pandas.DataFrame()
        
        if company_list != None:  
            self.company_list = self._get_csv(company_list)
            if self.company_list["Companies"].size > 0:
                self.stock_data = self.get_stock_data(self.company_list, \
                start_date, end_date)
        else:
            self.company_list = company_list
#        except TypeError, ex:
#            raise TypeError(ex.message)
            
    def _validate_dates(self, start_date, end):
        """Generate a list of days within the specified date and the number of days, or within the specified start and end date.

        """
        try:
            if isinstance(start_date, str):
                start_date = datetime.datetime.strptime(start_date, \
                "%Y-%m-%d").date()
            if isinstance(end, str):
                end = datetime.datetime.strptime(end, "%Y-%m-%d").date()
        except:
            raise ValueError("Date has a wrong format, try YYYY-MM-DD, e.g. 2013-01-01.")
        
        if not isinstance(start_date, datetime.date):
            TypeError("The dates has to either be a string or a date objects.")

        if isinstance(end, int):
            end = start_date + datetime.timedelta(days=end-1)        
        
        if end > datetime.datetime.now().date():
            raise ValueError("You are trying to get data from the future...")
            
        if (end - start_date).days < 0:
            raise ValueError("Number of days has to be greater than zero.")
        self.start_date = start_date
        self.end_date = end
        
    def _get_csv(self, file_source):
        """Read the csv file containing the company names from url."""
        try:
            if isinstance(file_source, str) and '.' in file_source:
                data = pandas.read_csv(file_source)
                if data.columns[0] != "Companies":
                    raise ValueError("The urls content is invalid.")
                return data
            elif isinstance(file_source, list):
                return pandas.DataFrame({"Companies" : file_source})
            else:
                return pandas.DataFrame({"Companies" : [file_source]})
        except urllib2.URLError, ex: # HTTPError is a subclass of URLError
            raise urllib2.URLError(ex.reason)
        except IOError, ex:
            raise IOError(ex.message)

    def _get_data(self):
        """Collect the stock data for specified companies. Save the closing
        stock prices for a specified period of time for a specified companies.
        Data stored in self.stock_data.

        """
        original_company_list = [i for i in self.company_list["Companies"]]
        temp_stock_data = []
        temp_company_list = []
        
        # datelist needed for initialization of the dictionary
        r = (self.end_date + datetime.timedelta(days=1) - self.start_date).days
        date_list = [self.start_date + datetime.timedelta(days=i) for i in xrange(r)]
                
        
        for single_company in original_company_list:
            try:
                quotes = quotes_historical_yahoo(single_company, self.start_date, self.end_date) # get the data from yahoo finance
                quote_dict = dict( [ (single_date, numpy.nan) for single_date in date_list ] ) # initialization of the dictionary with dates and nan values
                for row in quotes:
                    quote_dict[ num2date(row[0]).date() ] = row[1]
                # row[0] - date, row[1] - closing price
                temp_stock_data.append(quote_dict)
                temp_company_list.append(single_company)
            except:
                print 'Could not download data about: ' + single_company
        #Convert the fetched data to a pandas table.
        if(temp_stock_data == []): # if we did not fetch anything, just return
            return self.stock_data

        self.company_list = pandas.DataFrame({"Companies" : temp_company_list})
        temp = dict()
        for i in xrange(0, len(temp_company_list)):
            temp[original_company_list[i]] = temp_stock_data[i]
        stock_data = pandas.DataFrame(temp, index = temp[self.company_list["Companies"][0]].keys())
        
        stock_data = stock_data.dropna(how='all')
        stock_data = stock_data.sort()
        self.start_date = stock_data.index[0]
        self.end_date = stock_data.index[-1]
        return stock_data.combine_first(self.stock_data)
    
    
    def get_stock_data(self, companies, start_date, end):
        """Collect the stock data for specified company. Save the closing stock prices for a specified period of time for a specified companies. Data stored in self.stock_data.

		>>> data = sd.get_stock_data(["LNKD", "GOOG"], datetime.date(2013, 1, 1), 3)
		>>> data.columns.size
		2
		>>> data.index.size
		2
		
		As you can see in the example above the number of columns (so the number of companies) is, as expected, 2. However the number of rows (indexes) differ - instead of 3 dates it consist of two rows. It has been cause by the fact, that on the 1st of January there was no stock price. The package provides the removal of unnecessary dates, on which the stock prices for all companies would be equal to Nan.
		
        """
		
        try:
            self._validate_dates(start_date, end)
            if not isinstance(companies, pandas.DataFrame): 
                self.company_list = self._get_csv(companies)
            self.stock_data = self._get_data()
            return self.stock_data
        except TypeError, ex:
            raise TypeError(ex.message)

    def save_data(self, filename):
        """Save the data to a csv file."""
        try:
            self.stock_data.to_csv(filename)
        except TypeError, ex:
            raise TypeError(ex.message)

    def load_data(self, filename):
        """Load the data from a csv file.
		
		
		"""
        try:
            self.stock_data = pandas.read_csv(filename, index_col=0)
            date_list = [datetime.datetime.strptime(i, "%Y-%m-%d").date() for i in self.stock_data.index]
            self.start_date = date_list[0]
            self.end_date = date_list[-1]
            self.company_list = [i for i in self.stock_data.columns]
        except IOError, ex:
            raise IOError(ex.message)
            
    def get_stock_change(self, company, date, delta):
        """Returns the change of stock calculated by difference between average stock before this date and after this date. The delta defines how many days are taken into average calculation.

		>>> sd = stockdata.StockData()
		>>> sd.get_stock_change("LNKD", datetime.date(2013, 1, 2), 2)
		1.0096077215974495
		
		For the case that the company was not present at the stock and there are no stock prices for this company within a given timespan, the function returns 0. See LinkedIn example, when trying to fetch the data from year 2011.
		
		>>> sd.get_stock_change("LNKD", datetime.date(2011, 1, 2), 2)
		0
		
        """
        try:
            if delta <= 0:
                raise ValueError("Delta has to be greater than zero.")
    
            date = datetime.date(date.year, date.month, date.day)
            temp_time_delta = datetime.timedelta(days = delta + 5)
            time_delta = datetime.timedelta(days = delta)
    
            if company not in self.stock_data.keys():
                print "The data have not been fetched, fetching now."
                self.get_stock_data(company, date - temp_time_delta, date + temp_time_delta)
    
            if self.stock_data[company].index.searchsorted(date) < delta:
                print "The data have not been fetched, fetching now."
                self.get_stock_data(company, date - temp_time_delta, date)
    
            if abs(self.stock_data[company].index.searchsorted(date - time_delta) - len(self.stock_data[company].keys())) < delta:
                print "The data have not been fetched, fetching now."
                self.get_stock_data(company, date, date + temp_time_delta)
    
            curr_idx = self.stock_data[company].index.searchsorted(date)
    
            if (date + time_delta) < datetime.datetime.now().date():
                # before the date
                data_last_week = self.stock_data[company][curr_idx - delta:curr_idx]
                data_next_week = self.stock_data[company][curr_idx:curr_idx + delta]
    
                change = data_next_week.mean() / data_last_week.mean()
                if numpy.isnan(change):
                    change = 0
                return change
            return 0
        except TypeError, ex:
            raise TypeError(ex.message)
        except AttributeError, ex:
            raise AttributeError(ex.message)